package com.hoe.tsb.ring;

import com.hoe.tsb.fraft.exception.FRaftException;
import org.bouncycastle.asn1.x9.X962NamedCurves;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.jce.provider.JCEECPublicKey;
import org.bouncycastle.jce.provider.JDKKeyFactory;
import org.bouncycastle.math.ec.ECConstants;
import org.bouncycastle.math.ec.ECPoint;
import org.springframework.lang.NonNull;

import java.io.IOException;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.ECGenParameterSpec;
import java.util.Arrays;
import java.util.Base64;

/**
 * <p>
 * 环签名者
 * </p>
 *
 * @author hoe
 * @version 1.0
 * @date 2022/11/28 17:24
 */
public class RingSigner implements Signer {

    /**
     * 之所以定义这个数量级，是因为生成的128bit的随机si环恰好是这个<br/>
     * 如果修改了si的bit数，则需要对应调整这里的大小，使其取模后的大小看起来和生成的随机数值域相当
     */
    private static final BigInteger M = new BigInteger("999999999999999999999999999999999999999");
    /**
     * 标记字节数组的字节数
     */
    private static final int MARK_BYTES_LEN_NUM = 2;

    private final byte[] mPrivateKey;
    private final byte[][] mPublicKeys;
    private final int mRelateIdx;

    public RingSigner(String privateKey, String[] publicKeys, int relateIdx) {
        mPrivateKey = Base64.getDecoder().decode(privateKey);
        mPublicKeys = Arrays.stream(publicKeys).map(p -> Base64.getDecoder().decode(p)).toArray(byte[][] :: new);
        mRelateIdx = relateIdx;
    }

    @Override
    public byte[][] sign(byte[] data) {
        //生成G点
        ECPoint g = getStandardG();

        //生成随机数K
        //WARN 对于K值的管理可以是每次签名都生成，也可以是固定的
        BigInteger k = genK();

        //生成递推式的随机数环
        BigInteger[] si = genSiCircle();

        //计算k·g
        ECPoint kg = g.multiply(k);

        //计算C_i+1
        //C_i+1=Hash(m,S_i·g+C_i·pk_i)=Hash(m,k·g)
        BigInteger ci1 = eccHash(data, kg);

        //求出ci环
        //循环计算其他C
        //C_i=Hash(m,S_i-1·g+C_i-1·pk_i-1)
        BigInteger[] ci = new BigInteger[mPublicKeys.length];
        ci[mRelateIdx == mPublicKeys.length - 1 ? 0 : mRelateIdx + 1] = ci1;
        for (int cnt = 0,len = mPublicKeys.length, idx = mRelateIdx + 2; cnt < len - 1; cnt++, idx++) {
            if (idx >= len) {
                idx = 0;
            }

            //找出上一条的索引并最终找出si随机数
            int cIdx = idx == 0 ? len - 1 : idx - 1;

            //计算si·g
            ECPoint siG = g.multiply(si[cIdx]);
            //计算C_i-1·pk_i-1
            ECPoint pubPoint = convertPublicKeyToECPoint(mPublicKeys[cIdx]);
            ECPoint ciPki_1 = pubPoint.multiply(ci[cIdx]);
            //二者相加
            ECPoint preC_i_1 = siG.add(ciPki_1);
            //计算Hash
            ci[idx]=eccHash(data, preC_i_1);
        }

        //求出si
        BigInteger ciSki = ci1.multiply(new BigInteger(mPrivateKey));
        BigInteger sii = k.subtract(ciSki);
        si[mRelateIdx] = sii.mod(M);

        //构建签名并构造返回结果
        byte[][] result = new byte[2][];
        result[0] = genSign(ci[0], si);
        result[1] = k.toByteArray();
        return result;
    }

    @Override
    public boolean verify(byte[] data, byte[] sign) {
        //将ci1和si从sign中抽取出来
        int len = ((sign[0] & 0xFF) << 8) | (sign[1] & 0xFF);
        int extractIdx = MARK_BYTES_LEN_NUM;
        byte[] ci1Bytes = extractFrom(sign, extractIdx);
        extractIdx += MARK_BYTES_LEN_NUM + ci1Bytes.length;
        BigInteger ci0 = new BigInteger(ci1Bytes);
        BigInteger[] si = new BigInteger[len];
        byte[][] publicKeys = new byte[len][];
        for (int idx = 0; idx < len; idx++) {
            byte[] siItem = extractFrom(sign, extractIdx);
            extractIdx += MARK_BYTES_LEN_NUM + siItem.length;

            byte[] pbkItem = extractFrom(sign, extractIdx);
            extractIdx += MARK_BYTES_LEN_NUM + pbkItem.length;

            si[idx] = new BigInteger(siItem);
            publicKeys[idx] = pbkItem;
        }

        //生成G点
        ECPoint g = getStandardG();

        //验证签名，根据递推式依次求出ci环，看第一个是否和给出的一致
        //C_i+1=Hash(m,S_i·g+C_i·pk_i)=Hash(m,k·g)
        BigInteger[] cVrfArr = new BigInteger[len];
        cVrfArr[0] = ci0;
        for (int idx = 1, last = idx - 1; idx < len && idx != 0; idx++, last++) {
            if (idx > len - 1) {
                idx = 0;
                last = len - 1;
            }

            //计算si·g
            ECPoint siG = g.multiply(si[last]);
            ECPoint pubPoint = convertPublicKeyToECPoint(publicKeys[last]);
            //计算C_i·pk_i
            ECPoint ciPki_1 = pubPoint.multiply(cVrfArr[last]);
            //二者相加
            ECPoint hashParam = siG.add(ciPki_1);
            cVrfArr[idx]=eccHash(data, hashParam);
        }

        return ci0.equals(cVrfArr[0]);
    }

    /**
     * 生成签名<br/>
     * 前两字节表示公钥个数；其次两字节表示ci1字节数，其次是ci1的字节；之后是si和公钥字节，前面分别有两字节表示长度
     * @param ci1 第一个ci
     * @param si si环
     * @return 生成的签名二进制数组
     */
    private byte[] genSign(BigInteger ci1, BigInteger[] si) {
        //计算签名的长度
        //前2字节表示公钥个数
        int byteNumLen = MARK_BYTES_LEN_NUM;
        int totalLen = 2;

        byte[] ci1Byte = ci1.toByteArray();
        totalLen += byteNumLen + ci1Byte.length;
        byte[][] siBytes = new byte[si.length][];
        for (int idx = 0, len = si.length; idx < len; idx++) {
            siBytes[idx] = si[idx].toByteArray();
            totalLen += byteNumLen + siBytes[idx].length;
            totalLen += byteNumLen + mPublicKeys[idx].length;
        }

        //按照相应的算法将字节数组依次拷入
        int copyIndex = 0;
        byte[] result = new byte[totalLen];
        copyIndex = copyTo(result, shortToBytes(si.length), copyIndex, true);
        copyIndex = copyTo(result, ci1Byte, copyIndex);
        for (int idx = 0, len = si.length; idx < len; idx++) {
            copyIndex = copyTo(result, siBytes[idx], copyIndex);
            copyIndex = copyTo(result, mPublicKeys[idx], copyIndex);
        }

        return result;
    }

    /**
     * 将输入字节数组拷贝到结果字节数组中<br/>
     * 拷贝前会先计算输入数组的长度，并将长度用两字节进行记录放在首两字节
     * @param result 结果字节数组
     * @param input 输入字节数组
     * @param resultIdx 结果数组的开始拷贝下标
     * @return 拷贝到的结果数组下标
     */
    private int copyTo(@NonNull byte[] result, @NonNull byte[] input, int resultIdx) {
        return copyTo(result, input, resultIdx, false);
    }

    /**
     * 将输入字节数组拷贝到结果字节数组中<br/>
     * 拷贝前会先计算输入数组的长度，并将长度用两字节进行记录放在首两字节
     * @param result 结果字节数组
     * @param input 输入字节数组
     * @param resultIdx 结果数组的开始拷贝下标
     * @param onlyInput 是否仅拷贝数据本身
     * @return 拷贝到的结果数组下标
     */
    private int copyTo(@NonNull byte[] result, @NonNull byte[] input, int resultIdx, boolean onlyInput) {
        if (!onlyInput) {
            int len = input.length;
            byte[] lenBytes = shortToBytes(len);
            System.arraycopy(lenBytes, 0, result, resultIdx, lenBytes.length);
            resultIdx += lenBytes.length;
        }
        System.arraycopy(input, 0, result, resultIdx, input.length);
        return resultIdx + input.length;
    }

    /**
     * 将字节数组从签名中抽取出来
     * @param sign 环签名
     * @param extractIndex 开始抽取的下标
     * @return 抽取出来的字节数组
     */
    private byte[] extractFrom(@NonNull byte[] sign, int extractIndex) {
        int bytesLen = ((sign[extractIndex] & 0xFF) << 8) | (sign[extractIndex + 1] & 0xFF);
        byte[] result = new byte[bytesLen];
        System.arraycopy(sign, extractIndex + MARK_BYTES_LEN_NUM, result, 0, bytesLen);
        return result;
    }

    /**
     * 两字节整型转字节数组<br/>
     * java没有无符号short所以使用int替代
     * @param s 短整型
     * @return 字节数组
     */
    private byte[] shortToBytes(int s) {
        byte[] bytes = new byte[2];
        bytes[0] = (byte) ((s >> 8) & 0xFF);
        bytes[1] = (byte) (s & 0xFF);
        return bytes;
    }


    /**
     * 获取BigInteger的字节数
     * @param bi 待获取的大数
     * @return 字节数
     */
    private int getBigIntegerByteLen(BigInteger bi) {
        int bitLen = bi.bitLength();
        int fullLen = bitLen / 8;
        boolean hasTailLen = bitLen % 8 != 0;
        return fullLen + (hasTailLen ? 1 : 0);
    }

    /**
     * 获取标准G点<br/>
     * 实际签名时可以根据业务情况选取合适的G点
     * @return G点
     */
    private ECPoint getStandardG() {
        ECGenParameterSpec prime256v1 = new ECGenParameterSpec("prime256v1");
        String curveName = prime256v1.getName();

        X9ECParameters ecP = X962NamedCurves.getByName(curveName);
        if (ecP == null) {
            throw new RuntimeException("unknown curve name: " + curveName);
        }

        return ecP.getG();
    }

    /**
     * 生成标准的随机数K<br/>
     * 即256bit的随机数，也可以根据实际业务进行调整
     * @return 随机数K
     */
    private BigInteger genK() {
        return genRandom(256);
    }

    /**
     * 生成随机的递推式环<br/>
     * 除私钥对应的公钥外都要生成对应的随机数
     * @return 除指定公钥的递推式随机数环
     */
    private BigInteger[] genSiCircle() {
        BigInteger[] circleSi = new BigInteger[mPublicKeys.length];
        for (int idx = 0, len = mPublicKeys.length; idx < len; idx++){
            if (idx == mRelateIdx) {
                continue;
            }
            circleSi[idx] = genRandom(128);
        }
        return circleSi;
    }

    /**
     * 生成制定比特位的大数
     * @param numBits 比特位数
     * @return 生成后的大数
     */
    private BigInteger genRandom(int numBits) {
        BigInteger k;
        SecureRandom random = new SecureRandom();
        do {
            k = new BigInteger(numBits, random);
        }
        while (k.equals(ECConstants.ZERO));
        return k;
    }

    /**
     * 将ECC公钥转换为对应的ECPoint
     * @param publicKey 公钥
     * @return 转换后的ECPoint
     */
    private ECPoint convertPublicKeyToECPoint(byte[] publicKey) {
        try {
            PublicKey publicKeyFromDERStream = JDKKeyFactory.createPublicKeyFromDERStream(publicKey);
            return ((JCEECPublicKey) publicKeyFromDERStream).getQ();
        } catch (IOException e) {
            throw new FRaftException("将公钥转换为ECPoint错误", e);
        }
    }

    /**
     * 基于ECC进行摘要计算并最终映射到大数M的有限域上
     * @param data 待计算的数据
     * @param p 对应环的椭圆曲线点
     * @return ECC摘要
     */
    private BigInteger eccHash(byte[] data, ECPoint p) {
        byte[] digest;
        try {
            byte[] pBytes = (p.getX().toString() + p.getY().toString()).getBytes(StandardCharsets.UTF_8);
            byte[] mergeBytes = new byte[data.length + pBytes.length];
            System.arraycopy(data, 0, mergeBytes, 0, data.length);
            System.arraycopy(pBytes, 0, mergeBytes, data.length, pBytes.length);

            MessageDigest messageDigest = MessageDigest.getInstance("SHA-256");
            messageDigest.update(mergeBytes);
            digest = messageDigest.digest();
        } catch (Exception e) {
            throw new FRaftException("eccHash出错", e);
        }
        //映射到大数M上
        return new BigInteger(digest).mod(M);
    }

}
